package com.breeze.components.limiter.implementation;

import cn.hutool.cache.impl.LRUCache;
import com.breeze.components.core.exception.CommonException;
import com.breeze.components.limiter.anno.RateLimit;
import com.google.common.util.concurrent.RateLimiter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

/**
 * @author breeze
 */
@SuppressWarnings("UnstableApiUsage")
@RequiredArgsConstructor
@Slf4j
@Component
public class MapRateLimitChecker extends AbstractRateLimitChecker {

  /** 最大仅支持4096个key 超出这个key 限流将可能失效 */
  private final LRUCache<String, RateLimiter> cache = new LRUCache<>(4096);

  @Override
  public void check(RateLimit rateLimit, String combinedKey) {
    RateLimiter rateLimiter =
        cache.get(
            combinedKey, () -> RateLimiter.create((double) rateLimit.count() / rateLimit.time()));
    if (!rateLimiter.tryAcquire()) {
      throw new CommonException(rateLimit.message());
    }
    log.info("限制请求key:{}, combined key:{}", rateLimit.key(), combinedKey);
  }

  @Override
  public void tryAcquire(
      Runnable runnable, String combinedKey, int permits, int rateLimit, int rateTime) {
    RateLimiter rateLimiter =
        cache.get(combinedKey, () -> RateLimiter.create((double) rateLimit / rateTime));
    if (rateLimiter.tryAcquire(permits)) {
      runnable.run();
    }
  }
}
